from pylab import *
from camori_aux import step_renderer, step_get_colors

import scipy.optimize

import scipy.interpolate

if __name__ == '__main__':

    ion()


    # N = 2


    # # p = [1.0, -.5]
    # c1 = array([255,254,253])
    # #c2 = array([128,127,120])
    # c2 = array([20,10,0])


    # pL = [[-1.0,0.2],
    #       [0.0,0.0],
    #       [0.2,0.1],
    #       [2.0,-1.5]]    

    # for n,p in enumerate(pL):
    #     frmhat = zeros((2*N+1,2*N+1,3), dtype=np.int)
    #     step_renderer(p[0], p[1], c1, c2, frmhat)


    #     s = frmhat.sum(2)/3.0/255. - .5
    #     #h = outer(hamming(5), hamming(5))
    #     #s = s*h

    #     subplot(3,4,n+1)
    #     imshow(s, cmap=cm.bone, vmin = -.6, vmax=.6)

    #     S = fft(s)/10

    #     subplot(3,4,n+5)
    #     imshow(np.abs(S), cmap = cm.bone, vmin=-1, vmax=1)
    #     subplot(3,4,n+9)
    #     imshow(np.angle(S), cmap = cm.bone, vmin=-pi, vmax=pi)


    


    snp = np.load('snp.npz')['arr_0']
    snp = snp[15:-15, 15:-15]

    # snp = zeros((11,11))
    # for j in range(11):
    #     for k in range(11):
    #         if j<k:
    #             snp[j,k] = -60
    #         else:
    #             snp[j,k] = +60
                

    



    
    #h = outer(hamming(41), hamming(41))
    #h = ones(41)
    h = outer(ones(11), hamming(11))


    r = ones(11)
    #r[20]=0
    r[:5] = -1
    r = -r*60 * hamming(11)

    R = fft(r)

    snp = (snp-mean(snp))# * h

    imshow(snp, cmap=cm.bone)

    S = fft(snp*h)

    figure(2)
    subplot(2,1,1)
    plot(r, '-+')
    plot(snp[:5].T,'-+', label='1')
    legend()
    grid()
    subplot(2,1,2)
    plot(r, '-+')
    plot(snp[4:].T,'-+', label='1')
    legend()
    grid()

    ph = zeros((11,6))
    md = zeros((11,6))

    for N in range(10):
        # Z = S[20+N-5] * conj(S[20])
        Z = S[5+N-5] * conj(R)
        ph[N] = np.angle(Z[:6])


    RR = R * conj(R)
    rr = fftshift(np.angle(RR[:6]))

    figure(4)    
    subplot(2,1,1)
    plot(rr)
    for N in range(5):
        plot(ph[N], 'r-+')
    grid()
    subplot(2,1,2)
    plot(rr)

    xx = mgrid[:6]

    for N in range(5,11):
        plot(ph[N], 'b-+')
        plot(xx, -.4-.1*xx**2, 'k-x')
        plot(xx, -.4-.1*xx**2, 'k-x')

    grid()


    myx = ((4.5-ph[:,1] * 11 / (2*pi))[:10])

    yy = mgrid[:10]


    aa = np.polyfit(yy, myx, 1)

    print aa

    figure(5)
    imshow(snp, cmap=cm.bone)
    plot(myx, yy, 'g-+')
    plot( aa[0]*yy+aa[1],yy, 'r-+')




    # plot(4.5 - ph[:,1] * 11 / (2*pi), mgrid[:11], 'g-+')
    # plot(4.5 - ph[:,2] * 11 / (2*pi)/2., mgrid[:11], 'g-+')
    # plot(4.75 - ph[:,2] * 11 / (2*pi)/1.7, mgrid[:11], 'g-+')
